// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/common/base/simd_util.h>

#include <melon/benchmark.h>
#include <melon/init/init.h>

#include <random>

namespace kumo::pollux {
namespace {

#define POLLUX_BENCHMARK(_type, _name, ...) \
  [[maybe_unused]] _type _name(MELON_PP_STRINGIZE(_name), __VA_ARGS__)

template <typename T>
class LeadingMask {
 public:
  LeadingMask(const char* name, std::default_random_engine& gen) {
    std::uniform_int_distribution<> dist(0, ksimd::batch<T>::size + 1);
    for (int i = 0; i < kSize; ++i) {
      inputs_[i] = dist(gen);
    }
    melon::addBenchmark(__FILE__, name, [this] { return run(); });
  }

 private:
  unsigned run() {
    ksimd::batch_bool<T> ans = {};
    for (int i = 0; i < kSize; ++i) {
      ans = ans ^ simd::leadingMask<T>(inputs_[i]);
    }
    melon::doNotOptimizeAway(ans);
    return kSize;
  }

  static constexpr int kSize = 4 << 10;
  int8_t inputs_[kSize];
};

template <typename T>
class FromBitMask {
 public:
  FromBitMask(const char* name, std::default_random_engine& gen) {
    std::uniform_int_distribution<uint64_t> dist(
        0, (1ull << ksimd::batch<T>::size) - 1);
    for (int i = 0; i < kSize; ++i) {
      inputs_[i] = dist(gen);
    }
    melon::addBenchmark(__FILE__, name, [this] { return run(); });
  }

 private:
  unsigned run() {
    ksimd::batch_bool<T> ans = {};
    for (int i = 0; i < kSize; ++i) {
      ans = ans ^ simd::fromBitMask<T>(inputs_[i]);
    }
    melon::doNotOptimizeAway(ans);
    return kSize;
  }

  static constexpr int kSize = 2 << 10;
  uint64_t inputs_[kSize];
};

} // namespace
} // namespace kumo::pollux

int main(int argc, char* argv[]) {
  using namespace kumo::pollux;
  melon::Init melonInit(&argc, &argv);
  std::default_random_engine gen(std::random_device{}());
  POLLUX_BENCHMARK(LeadingMask<int32_t>, leadingMaskInt32, gen);
  POLLUX_BENCHMARK(LeadingMask<int64_t>, leadingMaskInt64, gen);
  POLLUX_BENCHMARK(FromBitMask<int32_t>, fromBitMaskInt32, gen);
  POLLUX_BENCHMARK(FromBitMask<int64_t>, fromBitMaskInt64, gen);
  melon::runBenchmarks();
  return 0;
}
